import argparse
import logging
import sys
from pathlib import Path
from sklearn.metrics import classification_report
from warmup_scheduler import GradualWarmupScheduler

import torch
import torch.nn as nn
import torch.nn.functional as F
import wandb
from torch import optim
from torch.utils.data import DataLoader, random_split, ConcatDataset
from tqdm import tqdm

from utils.data_loading import BasicDataset, CarvanaDataset
from utils.dice_score import dice_loss, CrossEntropyLoss_label_smooth
from evaluate import evaluate, evaluate_classification, evaluate_all_folders
from unet import UNet
from utils.transforms import ToTensor, Compose, RandomRotate, Center_Crop, RandomAffine
from utils.utils import seed_everything, kfolder

# # 1 channel
# dir_img = Path('./data/imgs/')
# dir_mask = Path('./data/masks/')

# 3 channel
dir_img = Path('./data/img/')
dir_mask = Path('./data/mask/')
dir_checkpoint = Path('./checkpoints/')

seed_everything(2022)


def train_net(folder_index,
              train_set, val_set,
              device,
              experiment,
              epochs: int = 5,
              batch_size: int = 1,
              learning_rate: float = 1e-5,
              val_percent: float = 0.1,
              save_checkpoint: bool = True,
              img_scale: float = 0.5,
              amp: bool = False,
              classification: bool = True):
    # 1.Change here to adapt to your data
    # n_channels=3 for RGB images
    # n_classes is the number of probabilities you want to get per pixel
    net = UNet(n_channels=4, n_classes=args.classes, bilinear=args.bilinear, device=device,classification =classification)
    # logging.info(f'Network:\n'
    #              f'\t{net.n_channels} input channels\n'
    #              f'\t{net.n_classes} output channels (classes)\n'
    #              f'\t{"Bilinear" if net.bilinear else "Transposed conv"} upscaling')
    if args.load:
        net.load_state_dict(torch.load(args.load, map_location=device))
        logging.info(f'Model loaded from {args.load}')

    net.to(device=device)
    # 2. Split into train / validation partitions
    # n_val = int(len(dataset) * val_percent)
    # n_train = len(dataset) - n_val
    # train_set, val_set = random_split(dataset, [n_train, n_val], generator=torch.Generator().manual_seed(0))
    val_set.transforms = Compose([
        Center_Crop(1024)
    ])

    # 3. Create data loaders
    loader_args = dict(batch_size=batch_size, num_workers=12, pin_memory=True)
    train_loader = DataLoader(train_set, shuffle=True, **loader_args)
    val_loader = DataLoader(val_set, shuffle=False, drop_last=False, **loader_args)

    logging.info(f'''Starting training:
        Epochs:          {epochs}
        Batch size:      {batch_size}
        Learning rate:   {learning_rate}
        Training size:   {len(train_set)}
        Validation size: {len(val_set)}
        Checkpoints:     {save_checkpoint}
        Device:          {device.type}
        Images scaling:  {img_scale}
        Mixed Precision: {amp}
    ''')

    # 4. Set up the optimizer, the loss, the learning rate scheduler and the loss scaling for AMP
    # optimizer = optim.RMSprop(net.parameters(), lr=learning_rate, weight_decay=1e-8, momentum=0.9)
    # scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=10)  # goal: maximize Dice score

    optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9)
    warmup_ratio = 100
    warmup_epoch = int(epochs * 0.2)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs - warmup_epoch, eta_min=0, last_epoch=-1)
    scheduler = GradualWarmupScheduler(optimizer, multiplier=warmup_ratio, total_epoch=warmup_epoch,
                                       after_scheduler=scheduler)

    grad_scaler = torch.cuda.amp.GradScaler(enabled=amp)
    criterion = nn.CrossEntropyLoss().to(device)
    # global_step = 0

    # 5. Begin training
    for epoch in range(1, epochs + 1):
        net.train()
        epoch_loss = 0
        with tqdm(total=len(train_set), desc=f'Epoch {epoch}/{epochs}', unit='img') as pbar:
            for batch in train_loader:
                images, true_masks, true_classes = batch['image'], batch['mask'], batch['class']

                # assert images.shape[1] == net.n_channels, \
                #     f'Network has been defined with {net.n_channels} input channels, ' \
                #     f'but loaded images have {images.shape[1]} channels. Please check that ' \
                #     'the images are loaded correctly.'

                images = images.to(device=device, dtype=torch.float32)
                true_masks = (true_masks / 255).to(device=device, dtype=torch.long)
                true_classes = true_classes.to(device=device, dtype=torch.long)

                with torch.cuda.amp.autocast(enabled=amp):
                    if classification:
                        masks_pred, classes_pred = net(images)
                        cls_loss = CrossEntropyLoss_label_smooth(classes_pred, true_classes)
                        # cls_loss = criterion(classes_pred, true_classes)
                        seg_loss = 0
                        # # 最后150轮只优化分类
                        # if epoch < epochs*0.7:
                        for channel_index in range(len(masks_pred)):
                            mask_pred = masks_pred[channel_index]
                            true_mask = true_masks[:, :, :, channel_index]
                            # seg_loss += criterion(mask_pred, true_mask)
                            seg_loss += dice_loss(F.softmax(mask_pred, dim=1).float(),
                                                  F.one_hot(true_mask, net.n_classes).permute(0, 3, 1, 2).float(),
                                                  multiclass=True)
                        # if epoch> epochs*0.7:
                        #     seg_loss = seg_loss * 0.001
                        loss = net.calculate_MTL_loss(cls_loss, seg_loss)
                        # loss = cls_loss + seg_loss
                    else:
                        loss = 0
                        masks_pred = net(images)
                        loss += dice_loss(F.softmax(masks_pred, dim=1).float(),
                                              F.one_hot(true_masks, net.n_classes).permute(0, 3, 1, 2).float(),
                                              multiclass=True)

                optimizer.zero_grad(set_to_none=True)
                grad_scaler.scale(loss).backward()
                grad_scaler.step(optimizer)
                grad_scaler.update()

                pbar.update(images.shape[0])
                # global_step += 1
                epoch_loss += loss.item()

                pbar.set_postfix(**{'loss (batch)': loss.item()})

        # Evaluation after 1 epoch
        if classification:
            prob_all, label_all, val_score, accuracy, auc, sensitivity, specificity = evaluate_classification(net,
                                                                                                              val_loader,
                                                                                                              device)
            logging.info(
                'Validation cls acc: {},AUC: {},sensitivity {},specificity {}'.format(accuracy, auc, sensitivity,
                                                                                      specificity))
            experiment.log(
                {f'folder{folder_index}/metrics/accuracy': accuracy, f'folder{folder_index}/metrics/AUC': auc,
                 f'folder{folder_index}/metrics/sensitivity': sensitivity,
                 f'folder{folder_index}/metrics/specificity': specificity})
        else:
            val_score = evaluate(net, val_loader, device)
        # scheduler.step(val_score)
        scheduler.step()

        # log
        if classification:
            experiment.log({
                f'folder{folder_index}/train seg loss': seg_loss.item(),
                f'folder{folder_index}/train cls loss': cls_loss.item(),
                f'folder{folder_index}/train total loss': loss.item(),
                f'folder{folder_index}/sigma_cls': net.sigma_cls.item(),
                f'folder{folder_index}/sigma_seg': net.sigma_seg.item(),
            }, step=epoch + epochs * folder_index)
        else:
            experiment.log({
                f'folder{folder_index}/train loss': loss.item(),
            }, step=epoch + epochs * folder_index)
        logging.info('Validation seg Dice score: {}'.format(val_score))
        experiment.log({
            f'folder{folder_index}/learning rate': optimizer.param_groups[0]['lr'],
            f'folder{folder_index}/metrics/validation Dice': val_score,
        }, step=epoch + epochs * folder_index)
        # if epoch % 5 ==0:
        #     histograms = {}
        #     for tag, value in net.named_parameters():
        #         tag = tag.replace('/', '.')
        #         histograms['Weights/' + tag] = wandb.Histogram(value.data.cpu())
        #         histograms['Gradients/' + tag] = wandb.Histogram(value.grad.data.cpu())
        #     experiment.log({
        #     'images': wandb.Image(images[0].cpu()),
        #     'masks': {
        #         'true0': wandb.Image(true_masks[0,:,:,0].float().cpu()),
        #         'true1': wandb.Image(true_masks[0,:,:,1].float().cpu()),
        #         'true2': wandb.Image(true_masks[0,:,:,2].float().cpu()),
        #         'pred0': wandb.Image(masks_pred[0].argmax(dim=1)[0].float().cpu()),
        #         'pred1': wandb.Image(masks_pred[1].argmax(dim=1)[0].float().cpu()),
        #         'pred2': wandb.Image(masks_pred[2].argmax(dim=1)[0].float().cpu()),
        #     },
        #     **histograms
        #     },step=epoch)
    if save_checkpoint:
        Path(dir_checkpoint).mkdir(parents=True, exist_ok=True)
        torch.save(net.state_dict(), str(dir_checkpoint / 'folder{}checkpoint_epoch{}.pth'.format(folder_index, epoch)))
        logging.info(f'folder{folder_index} Checkpoint {epoch} saved!')

    return prob_all, label_all, val_score


def get_args():
    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
    parser.add_argument('--epochs', '-e', metavar='E', type=int, default=500, help='Number of epochs')
    parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=4, help='Batch size')
    parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-4,
                        help='Learning rate', dest='lr')
    parser.add_argument('--load', '-f', type=str, default=False, help='Load model from a .pth file')
    parser.add_argument('--scale', '-s', type=float, default=1.0, help='Downscaling factor of the images')
    parser.add_argument('--validation', '-v', dest='val', type=float, default=10.0,
                        help='Percent of the data that is used as validation (0-100)')
    parser.add_argument('--amp', action='store_true', default=False, help='Use mixed precision')
    parser.add_argument('--bilinear', action='store_true', default=False, help='Use bilinear upsampling')
    parser.add_argument('--classes', '-c', type=int, default=2, help='Number of classes')
    parser.add_argument('--classification', action='store_true', default=False, help='cls or seg')

    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()

    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')

    # (Initialize logging)
    experiment = wandb.init(project='U-Net', resume='allow', anonymous='must')
    experiment.config.update(dict(epochs=args.epochs, batch_size=args.batch_size, learning_rate=args.lr, ))

    prob_all_folder = []
    label_all_folder = []
    val_scores = 0

    # prepare kfolder dataset
    train_image_path_lists, test_image_path_lists = kfolder(num_folder=10, image_dir=dir_img, seed=2022,
                                                            shuffle=False)
    for folder_index in range(10):
        # Create dataset
        train_set = BasicDataset(dir_img, dir_mask, args.scale)
        val_set = BasicDataset(dir_img, dir_mask, args.scale)
        train_set.ids = train_image_path_lists[folder_index]
        val_set.ids = test_image_path_lists[folder_index]
        logging.info(
            f'Creating train dataset with {len(train_set)} examples,test dataset with {len(val_set)} examples.')

        # train
        prob_per_folder, label_per_folder, val_score = train_net(folder_index=folder_index,
                                                                 train_set=train_set, val_set=val_set,
                                                                 experiment=experiment,
                                                                 epochs=args.epochs,
                                                                 batch_size=args.batch_size,
                                                                 learning_rate=args.lr,
                                                                 device=device,
                                                                 img_scale=args.scale,
                                                                 val_percent=args.val / 100,
                                                                 amp=args.amp,
                                                                 classification=args.classification)
        prob_all_folder.extend(prob_per_folder)
        label_all_folder.extend(label_per_folder)
        val_scores += val_score
        evaluate_all_folders(prob_all_folder, label_all_folder, val_scores, experiment,
                             step=args.epochs * (folder_index + 1))
